from torch.utils.data import Dataset
from csv import reader as CSV_reader
from os.path import join as o_p_join
from PIL import Image
from torchvision.transforms import Compose, ToTensor, Normalize


CNAME_FILE = 'CLASS NAMES.csv'
ITEM_FILE = 'butterflies.csv'
NUM_CLASSES = 75


class BF75_Dset(Dataset):
    def __init__(self, root_dir, dstype):
        super(BF75_Dset, self).__init__()
        self.root_dir = root_dir
        # 创建类型列表
        self.classes = []
        with open(o_p_join(self.root_dir, CNAME_FILE)) as csv_f:
            csv_r = CSV_reader(csv_f)
            next(csv_r)
            for __ in range(NUM_CLASSES):
                row_data = next(csv_r)
                self.classes.append(row_data[0])
        # 创建类型字典
        self.class_dict = dict(zip(self.classes, list(range(NUM_CLASSES))))
        # 统计各项数据列表
        self.item_dir = []
        self.item_class = []
        self.item_count = 0
        with open(o_p_join(self.root_dir, ITEM_FILE)) as csv_f:
            csv_r = CSV_reader(csv_f)
            next(csv_r)
            while True:
                try:
                    row_data = next(csv_r)
                except StopIteration:
                    break
                if row_data[2] == dstype:
                    self.item_count += 1
                    self.item_dir.append(row_data[0])
                    self.item_class.append(self.class_dict[row_data[1]])
        # 确定数据预处理方式
        self.item_trans = Compose(
            [ToTensor(), Normalize([0.5, 0.5, 0.5], [0.2, 0.2, 0.2])])

    def __getitem__(self, idx):
        img_dir = o_p_join(self.root_dir, self.item_dir[idx])
        img = self.item_trans(Image.open(img_dir))
        return img, self.item_class[idx]

    def __len__(self):
        return self.item_count
